#!/usr/bin/python
# Copyright (c) 2010-, WangFei, Cloud Team, CGCL,
# All rights reserved.
# This program configures multiple virtual machines(VM) as a MPI environment.
# Using format is as follow:
#    python setmpienv --machinefile machine
# Machine file includes list of IP and Root'Password of VM. 
# And the first line of machine will as the Master Node. 
# This program is only for specific virtual machine.
# 
#
import sys
import os
import pexpect
import time
import pp
import paramiko

def set_etc_hosts(host_address,username,password,ssh_port,all_host_list):   
    cmd = "echo \'127.0.0.1       localhost.localdomain localhost\' > /etc/hosts;"
    for aip in all_host_list:
        cmd += "echo \'"+aip+" "+aip+"\' >> /etc/hosts;"
    cmd += "/bin/hostname "+host_address +";"     
    cmd += " echo '' > /etc/hosts.equiv;"
    for ip1 in all_host_list:
        cmd += "echo \'" +ip1 + "\' >> /etc/hosts.equiv;"
    cmd += "echo '' > /root/.rhosts;"
    for ip2 in all_host_list:
        cmd += "echo \'" +ip2 + "\' >> /root/.rhosts;"
    HOSTNAME=host_address
    NETWORKING="yes"
    NETWORKING_IPV6="no"
    cmd += " echo \'" + "NETWORKING="+NETWORKING+"\' > /etc/sysconfig/network;" \
            + "echo \'NETWORKING_IPV6=" + NETWORKING_IPV6 + "\' >> /etc/sysconfig/network ;" \
            + "echo \'HOSTNAME=" + HOSTNAME + "\' >> /etc/sysconfig/network; "
    client=paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((host_address), username=(username), password=(password), port=(ssh_port))
    stdin,stdout,stderr = client.exec_command(cmd)
    print stdout.read()
    client.close()

def set_salve_nfs_rsh(Salve_IP,Master_IP,username,password,ssh_port):
    cmd = "sed -i \'7d\' /etc/fstab;" + \
         "echo \'"+ Master_IP+":/home /home nfs defaults 0 0\' >> /etc/fstab ;" + \
         "/sbin/chkconfig netfs on;" + \
                "mount /home;" + \
		        "/sbin/chkconfig rsh on;" + \
                "/etc/init.d/xinetd start;"    
    client=paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((Salve_IP), username=(username), password=(password), port=(ssh_port))
    stdin,stdout,stderr = client.exec_command(cmd)
    print stdout.read()
    client.close()

def set_master_rsh(Master_IP,username,password,ssh_port):
    cmd = "/sbin/chkconfig rsh on;" + \
              "/etc/init.d/xinetd start;"
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((Master_IP), username=(username), password=(password), port=(ssh_port))
    stdin,stdout,stderr = client.exec_command(cmd)
    print stdout.read()
    client.close()

def set_master_nfs(Master_IP,username,password,ssh_port):
    cmd = "/sbin/chkconfig nfs on;" + \
        	"/sbin/chkconfig nfslock on;" + \
        	"/etc/init.d/nfslock start;" + \
        	"/etc/init.d/nfs start"
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((Master_IP), username=(username), password=(password), port=(ssh_port))
    stdin,stdout,stderr = client.exec_command(cmd)
    print stdout.read()
    client.close()	

def set_mpd_port(Master_IP, total_ip, username, password, ssh_port, all_host_list):
    cmd = "echo '' > /root/mpd.hosts;"
    for ip in all_host_list:
        cmd += "echo "+ip+" >> /root/mpd.hosts;"
    cmd += "mpdallexit; mpdboot --rsh=rsh -n "+str(total_ip)+" --file=/root/mpd.hosts"
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((Master_IP), username=(username), password=(password), port=(ssh_port))
    stdin,stdout,stderr = client.exec_command(cmd)
    print stdout.read()
    client.close()

def GetIpAndPassword(machinefile):
    global all_host_list
    global password_list
    global total_ip
    all_host_list=[]
    password_list=[]
    path=sys.path[0]
    fp_intr = open(path+'/'+machinefile,'r')
    total_ip = 0
    while True:
        line = fp_intr.readline()
        length = len(line)
        if length == 0: # Zero length indicates EOF
            break
        line=line[0:length-1] 
        ip,password=line.split(' ')
        #print ip
        #print password
        total_ip += 1
        all_host_list.append(ip)
        password_list.append(password)
    fp_intr.close()

def set_mpi_env():
    global all_host_list
    global total_ip
    global password_list
    ssh_port=22
    username="root"
    Master_IP=all_host_list[0]    
    start_time_total = time.time()
    start_time = time.time()
    ppservers = ()
    job_server = pp.Server(total_ip, ppservers=ppservers)
    print "Starting pp with", job_server.get_ncpus(), "workers"
    jobs = []
    index = 0
    while index < total_ip:
        host_address=all_host_list[index]
        password = password_list[index]
        jobs.append(job_server.submit(set_etc_hosts, (host_address,username,password,ssh_port,all_host_list,), (),("paramiko",)))
        index += 1

    for job in jobs:
        job()
    
    print "Time elapsed: ", time.time() - start_time, "s"
    job_server.print_stats()
    start_time = time.time()
    ppservers = ()
    job_server = pp.Server(2, ppservers=ppservers)
    print "Starting pp with", job_server.get_ncpus(), "workers"
    
    password=password_list[0]
    jobs4 = []
    jobs4.append(job_server.submit(set_master_nfs, (Master_IP,username,password,ssh_port,), (), ("paramiko",)))
    jobs4.append(job_server.submit(set_master_rsh, (Master_IP,username,password,ssh_port,), (), ("paramiko",)))
    for job in jobs4:
        job()
    print "Time elapsed: ", time.time() - start_time, "s"
    job_server.print_stats()
    start_time = time.time()
    ppservers = ()
    job_server = pp.Server(total_ip-1, ppservers=ppservers)
    print "Starting pp with", job_server.get_ncpus(), "workers"

    jobs1 = []
    index = 1
    while index < total_ip:
        host_address=all_host_list[index]
        password=password_list[index]
        jobs1.append(job_server.submit(set_salve_nfs_rsh, (host_address,Master_IP,username,password,ssh_port,), (), ("paramiko",)))
        index += 1 
    for job in jobs1:
        job()
    print "Time elapsed: ", time.time() - start_time, "s"
    job_server.print_stats()

    password=password_list[0]
    set_mpd_port(Master_IP,total_ip,username,password,ssh_port,all_host_list)
    print "Time elapsed: ", time.time() - start_time_total, "s"

all_host_list=[]
password_list=[]
total_ip=0
if __name__ == "__main__":
    cnt_argv = len(sys.argv)
    if cnt_argv != 3 or sys.argv[1] != "--machinefile":
        print "Command format error. Please Using As follow:\n"
        print "      python setmpienv --machinefile ip"
        sys.exit(0)
    machinefile=sys.argv[2]
    flag = os.path.exists(sys.path[0]+"/" + machinefile)
    if not flag:
        print "File "+machinefile+" is not exist."
        sys.exit(0)
    print "Machine list file is " + machinefile
    GetIpAndPassword(machinefile)
    set_mpi_env()


